package com.ribbon.ribbonClient.rule;

import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;

import java.util.List;
import java.util.Objects;

/**
 * 实现IRule接口，自定义负载规则
 * Created by earl on 2018/4/25.
 */
public class MyRule implements IRule{

    ILoadBalancer lb;
    public MyRule(){}

    public  MyRule(ILoadBalancer lb){
        this.lb = lb;
    }

    @Override
    public Server choose(Object key){
        int port =  Objects.isNull(key)?0: (int) key;
        //获取所有服务器
        List<Server> serverList = lb.getAllServers();
        //返回相同端口的服务器，若是没有，直接返回服务器列表的第一个服务器
        return serverList.stream().filter(server -> Objects.equals(server.getPort(), port)).findAny().orElse(serverList.get(0));
    }

    @Override
    public void setLoadBalancer(ILoadBalancer iLoadBalancer) {
        this.lb = iLoadBalancer;
    }

    @Override
    public ILoadBalancer getLoadBalancer() {
        return lb;
    }

}
